import cv2
from ultralytics import YOLO
from utils import img_utils
# Load a model
model = YOLO(r"C:\Users\kang_\Downloads\best (1).pt")  # load a pretrained model (recommended for training)
# Train the model
# cv2读取视频
cap = cv2.VideoCapture(r"C:\Users\kang_\Downloads\2月17日\2月17日(1).mp4")
cv2.namedWindow('w', cv2.WINDOW_NORMAL)
cv2.resizeWindow('w', 1920, 1080)
if cap.isOpened():
    while True:
        ret, frame = cap.read()
        # 如果读取失败或到达视频末尾，退出循环
        if not ret:
            break
        # frame = img_utils.light_balance(frame)
        results = model(frame, imgsz=640, visualize=False, conf=0.5)
        for result in results:
            cv2.imshow('w', result.plot())
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
